from pytorch_lightning.callbacks import ModelCheckpoint
import time
from pytorch_lightning import cli_lightning_logo, LightningDataModule, LightningModule
from pytorch_lightning.cli import LightningCLI
from datam import *
from model import *
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning import Trainer
from config import  config
from pytorch_lightning.callbacks import EarlyStopping




if __name__ == '__main__':

    pl.seed_everything(1234)
    checkpoint_callback = ModelCheckpoint(
        monitor='val_loss',
        dirpath=config.modelPATH,
        filename='sample-mnist原始-{epoch:02d}-{val_loss:.2f}',
        save_top_k=1,
        mode='min',
        save_last=True
    )
    early_stopping = EarlyStopping('val_loss')
    trainer = Trainer(gpus=config.AVAIL_GPUS, max_epochs=config.max_epochs, callbacks=[checkpoint_callback,early_stopping])
    data_mnist = DataM(config.data_dir,config.BATCH_SIZE,config.AVAIL_GPUS)
    data_mnist.setup()
    model = Model(Backbone())
    #训练模型
    trainer.fit(model,data_mnist)
    trainer.test(model,data_mnist.test_dataloader())





